import os
import pandas as pd
import numpy as np
import glob
import math
from pathlib import Path
from shutil import copyfile, move


def img_copy(img_path, img_name, code, dst_path, with_label_file, label_suffix):
    if with_label_file:
        label_path = Path(img_path).with_suffix(label_suffix)
        label_dst_path = os.path.join(dst_path, code, label_path.name)
        if os.path.exists(label_path):
            copyfile(str(label_path), label_dst_path)
            copyfile(img_path, os.path.join(dst_path, code, img_name))
    else:
        copyfile(img_path, os.path.join(dst_path, code, img_name))

    

def main():
    ps = [r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t9_DFO/T9_DFO/t9_DFO_defect/data_0517', 
        r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t9_DFO/T9_DFO/t9_DFO_defect/data_0418']
    dst1 = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t9_DFO/T9_DFO/t9_DFO_defect/train_0518'
    dst2 = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t9_DFO/T9_DFO/t9_DFO_defect/test_0518'

    limit_codes = []
    limit_nums = {'TGXID': 150}
    no_defect_code = 'TGXID'
    label_suffix = '.json'
    test_ratio = 0.15

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    df = pd.DataFrame()
    for path in ps:
        imgs = glob.glob(os.path.join(path, '*/*.jpg'))
        tmp_df = pd.DataFrame()
        tmp_df['img_path'] = imgs
        tmp_df['img_name'] = tmp_df['img_path'].parallel_apply(lambda x: Path(x).name)
        tmp_df['code'] = tmp_df['img_path'].parallel_apply(lambda x: Path(x).parent.name)
        df = pd.concat([df, tmp_df])
    
    print()
    print(df['code'].value_counts())

    if len(limit_codes) > 0:
        df = df[df['code'].isin(limit_codes)].reset_index(drop=True)
    
    for code in df['code'].unique():
        if code in limit_nums:
            df1 = df[df['code']==code]
            if(len(df1) <= limit_nums[code]):
                continue
            df2 = df[df['code']!=code]
            df1 = df1.sample(frac=1, random_state=10).reset_index(drop=True)
            df1 = df1.iloc[:limit_nums[code], :]
            df = pd.concat([df1, df2]).reset_index(drop=True)


    for code in df['code'].unique():
        
        tmp_df = df[df['code']==code].reset_index(drop=True)
        print('\ncode: %s, len: %d .'% (code, len(tmp_df)))

        train_num = math.ceil(len(tmp_df) * (1-test_ratio))
        os.makedirs(os.path.join(dst1, code), exist_ok=True)
        tmp_df.iloc[:train_num, :].parallel_apply(lambda x: img_copy(x['img_path'], x['img_name'], x['code'], dst1, x['code']!=no_defect_code, label_suffix), axis=1)

        if train_num < len(tmp_df):
            os.makedirs(os.path.join(dst2, code), exist_ok=True)
            tmp_df.iloc[train_num:, :].parallel_apply(lambda x: img_copy(x['img_path'], x['img_name'], x['code'], dst2, x['code']!=no_defect_code, label_suffix), axis=1)



if __name__ == '__main__':
    main()
